# Analytic inversion module for FLEXINVERT.jl
# Ports mod_analytic.f90 functionality for direct matrix solution

module AnalyticInversion

using LinearAlgebra
using ..PriorCovariance: CovarianceMatrix
using ..Transformations
using ..StateVector: InversionState
using ..ObservationsCore: Observations

export solve_analytic, AnalyticResult

"""
    AnalyticResult

Result structure for analytic inversion.
"""
struct AnalyticResult
    x::Vector{Float64}        # Solution in control space
    success::Bool             # Success flag
    posterior_covariance::Union{Matrix{Float64}, Nothing}  # Posterior covariance
    analysis_variance::Float64    # Total posterior variance
    dof::Float64             # Degrees of freedom for signal
end

"""
    solve_analytic(H::Matrix{Float64}, y::Vector{Float64}, R_inv::Matrix{Float64},
                  covariance::CovarianceMatrix) -> AnalyticResult

Solve the linear Bayesian inversion analytically using the normal equations.

The solution is: x_a = B * H^T * (H * B * H^T + R)^{-1} * (y - H * x_b)

In control space with phi = B^{-1/2} * (x - x_b):
phi_a = (I + U^T * H^T * R^{-1} * H * U * Λ)^{-1} * U^T * H^T * R^{-1} * (y - H * x_b)

# Arguments
- `H`: Observation operator (linearized Jacobian)
- `y`: Observations vector
- `R_inv`: Inverse observation error covariance matrix
- `covariance`: Prior covariance matrix with eigenvalue decomposition

# Returns
- AnalyticResult with solution and diagnostics
"""
function solve_analytic(H::Matrix{Float64}, y::Vector{Float64}, R_inv::Matrix{Float64},
                       covariance::CovarianceMatrix)

    m, n = size(H)  # m observations, n state variables

    if n != covariance.nx * covariance.ny * covariance.nt
        error("Observation operator dimension mismatch with state space")
    end

    @info "Solving analytic inversion: $m observations, $n state variables, $(covariance.n_modes) modes"

    # Project observation operator to control space via B^{1/2}
    H_reduced = Transformations.build_H_reduced(H, covariance)

    # Form the matrix to invert: I + H_reduced^T * R_inv * H_reduced
    # This is the control space Hessian of the cost function
    A = I + H_reduced' * R_inv * H_reduced
    A_sym = Symmetric(Matrix(A))

    # Right-hand side: H_reduced^T * R_inv * y
    rhs = H_reduced' * R_inv * y

    phi_solution = zeros(Float64, size(A, 1))
    posterior_cov_control = Matrix{Float64}(undef, size(A, 1), size(A, 1))
    used_cholesky = false
    truncated_modes = 0

    try
        chol = cholesky(A_sym; check=false)
        if chol.info == 0
            phi_solution = chol \ rhs
            posterior_cov_control .= inv(chol)
            used_cholesky = true
        else
            error("Cholesky factorization failed with info=$(chol.info)")
        end
    catch err
        eig = eigen(A_sym)
        max_eig = maximum(eig.values)
        tol = max(max_eig * eps(Float64) * max(1, size(A_sym, 1)) * 10, eps(Float64))
        inv_vals = zeros(Float64, length(eig.values))
        for j in eachindex(eig.values)
            if eig.values[j] > tol
                inv_vals[j] = 1.0 / eig.values[j]
            else
                inv_vals[j] = 0.0
                truncated_modes += 1
            end
        end

        if all(iszero, inv_vals)
            @error "Analytic inversion failed: control-space Hessian is singular with no resolvable modes" err
            return AnalyticResult(zeros(covariance.n_modes), false, nothing, 0.0, 0.0)
        end

        coeffs = inv_vals .* (eig.vectors' * rhs)
        phi_solution = eig.vectors * coeffs
        posterior_cov_control .= eig.vectors * Diagonal(inv_vals) * eig.vectors'

        if truncated_modes > 0
            @warn "Control-space Hessian is rank-deficient; truncated $truncated_modes modes" tolerance=tol
        end
    end

    posterior_cov_control = Symmetric(posterior_cov_control)

    # Transform back to physical space if needed
    posterior_cov_physical = nothing
    if covariance.n_modes <= 1000
        try
            B_full = Transformations.full_covariance_matrix(covariance)
            posterior_cov_physical = B_full - B_full * H' * inv(H * B_full * H' + inv(R_inv)) * H * B_full
        catch err
            @warn "Failed to compute posterior covariance in physical space" err
            posterior_cov_physical = nothing
        end
    end

    # Compute diagnostics
    analysis_variance = tr(posterior_cov_control)
    dof = tr(posterior_cov_control * (I - posterior_cov_control))

    if used_cholesky
        @info "Analytic solution computed successfully via Cholesky factorization"
    else
        @info "Analytic solution computed with eigenmode truncation"
    end
    @info "Analysis variance: $analysis_variance"
    @info "Degrees of freedom for signal: $dof"

    return AnalyticResult(phi_solution, true, posterior_cov_physical, analysis_variance, dof)
end

"""
    solve_analytic_simple(obs::Observations, H::Matrix{Float64}, covariance::CovarianceMatrix) -> AnalyticResult

Simplified analytic inversion with diagonal observation error covariance.

# Arguments
- `obs`: Observations with measurement errors
- `H`: Observation operator
- `covariance`: Prior covariance matrix

# Returns
- AnalyticResult with solution
"""
function solve_analytic_simple(obs::Observations, H::Matrix{Float64}, covariance::CovarianceMatrix)

    # Diagonal observation error covariance
    R_inv = Diagonal(1.0 ./ obs.measurement_errors.^2)

    # Observations vector
    y = obs.concentrations

    return solve_analytic(H, y, R_inv, covariance)
end

"""
    posterior_uncertainty(result::AnalyticResult, covariance::CovarianceMatrix) -> Vector{Float64}

Compute posterior uncertainty (standard deviations) for each state variable.

# Arguments
- `result`: Analytic inversion result
- `covariance`: Prior covariance matrix

# Returns
- Vector of posterior standard deviations
"""
function posterior_uncertainty(result::AnalyticResult, covariance::CovarianceMatrix)
    if !result.success || result.posterior_covariance === nothing
        return zeros(covariance.nx * covariance.ny * covariance.nt)
    end

    return sqrt.(diag(result.posterior_covariance))
end

"""
    compute_averaging_kernel(H::Matrix{Float64}, covariance::CovarianceMatrix, R_inv::Matrix{Float64}) -> Matrix{Float64}

Compute the averaging kernel matrix A = ∂x_a/∂x_true.

The averaging kernel shows how the retrieved state relates to the true state:
x_a = A * x_true + (I - A) * x_b + error

# Arguments
- `H`: Observation operator
- `covariance`: Prior covariance matrix
- `R_inv`: Inverse observation error covariance

# Returns
- Averaging kernel matrix
"""
function compute_averaging_kernel(H::Matrix{Float64}, covariance::CovarianceMatrix, R_inv::Matrix{Float64})

    # In physical space: A = B * H^T * (H * B * H^T + R)^{-1} * H
    B_full = Transformations.full_covariance_matrix(covariance)

    # Gain matrix
    K = B_full * H' * inv(H * B_full * H' + inv(R_inv))

    # Averaging kernel
    A = K * H

    return A
end

"""
    information_content(averaging_kernel::Matrix{Float64}) -> Float64

Compute the information content (degrees of freedom for signal) from averaging kernel.

Information content = trace(A) where A is the averaging kernel.

# Arguments
- `averaging_kernel`: Averaging kernel matrix

# Returns
- Information content (scalar)
"""
function information_content(averaging_kernel::Matrix{Float64})
    return tr(averaging_kernel)
end

"""
    validate_analytic_solution(result::AnalyticResult, H::Matrix{Float64}, y::Vector{Float64},
                              R_inv::Matrix{Float64}, covariance::CovarianceMatrix) -> Bool

Validate the analytic solution by checking residuals and cost function.

# Arguments
- `result`: Analytic solution result
- `H`: Observation operator
- `y`: Observations
- `R_inv`: Inverse observation error covariance
- `covariance`: Prior covariance matrix

# Returns
- `true` if solution is valid, `false` otherwise
"""
function validate_analytic_solution(result::AnalyticResult, H::Matrix{Float64}, y::Vector{Float64},
                                   R_inv::Matrix{Float64}, covariance::CovarianceMatrix)

    if !result.success
        return false
    end

    # Transform solution to physical space
    x_physical = Transformations.phi2chi(result.x, covariance)

    # Check observation residuals
    y_mod = H * x_physical
    residual = y - y_mod
    chi2_obs = residual' * R_inv * residual

    # Check if residuals are reasonable
    expected_chi2 = length(y)  # Expected value for chi-squared with length(y) DOF

    @info "Validation: χ² = $chi2_obs, expected ≈ $expected_chi2"

    # Solution is valid if chi-squared is not too large
    return chi2_obs < 5 * expected_chi2
end

end # module AnalyticInversion
